/*
 *   Copyright (c) 2020 HiSilicon (Shanghai) Technologies CO., LIMITED.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * Description: WAPI Implementation
 */

#include "securec.h"
#include "wapi.h"
#include "mbedtls/md.h"
#include "mbedtls/cipher.h"
#include "sha256.h"
#include "cipher.h"
#include "km.h"
#include "wpa_debug.h"

/* Maximum length of data generated by the HMAC_SHA256 operation */
#define HMAC_SHA256_SIZE 32
#ifndef BITS_PER_BYTE
#define BITS_PER_BYTE 8
#endif

int wai_hmac_sha256(const unsigned char *data, unsigned int datalen,
    const unsigned char *key, unsigned int keylen, unsigned char *output, unsigned int outlen)
{
    unsigned char hash[HMAC_SHA256_SIZE] = { 0 };

    if (data == NULL || key == NULL || output == NULL) {
        wpa_printf(MSG_WARNING, "The param for HMAC-SHA256 is wrong");
        return WAPI_FAILED;
    }

    /* HMAC-SHA256 support not more than HMAC_SHA256_SIZE */
    if ((outlen > HMAC_SHA256_SIZE) || (outlen == 0)) {
        wpa_printf(MSG_WARNING, "The length of HMAC-SHA256 dest is too long");
        return WAPI_FAILED;
    }

    /* the final will always be 32 Bytes */
    if (hmac_sha256(key, keylen, data, datalen, hash) < 0) {
        forced_memzero(hash, HMAC_SHA256_SIZE);
        wpa_printf(MSG_WARNING, "hmac_sha256 fail");
        return WAPI_FAILED;
    }
    /* Copy expected length */
    if (memcpy_s(output, outlen, hash, outlen) != EOK) {
        forced_memzero(hash, HMAC_SHA256_SIZE);
        wpa_printf(MSG_WARNING, "memcpy fail");
        return WAPI_FAILED;
    }

    forced_memzero(hash, HMAC_SHA256_SIZE);
    return WAPI_SUCCESS;
}

/**
 *  @brief KD-HMAC-SHA256 support any length output.
 */
int wai_kd_hmac_sha256(const unsigned char *data, unsigned int datalen,
    const unsigned char *key, unsigned int keylen, unsigned char *output, unsigned int outlen)
{
    int ret = WAPI_FAILED;
    unsigned int outidx;

    if (data == NULL || key == NULL || output == NULL) {
        wpa_printf(MSG_WARNING, "The param for HMAC-SHA256 is wrong");
        return WAPI_FAILED;
    }

    for (outidx = 0; outlen / HMAC_SHA256_SIZE; outidx += HMAC_SHA256_SIZE) {
        if ((ret = hmac_sha256(key, keylen, data, datalen, &output[outidx])) < 0) {
            wpa_printf(MSG_WARNING, "hmac_sha256 fail %d", ret);
            return WAPI_FAILED;
        }
        data = &output[outidx];
        datalen = HMAC_SHA256_SIZE;
        outlen -= HMAC_SHA256_SIZE;
    }

    if (outlen > 0) {
        ret = wai_hmac_sha256(data, datalen, key, keylen,
            &output[outidx], outlen);
        if (ret != WAPI_SUCCESS) {
            return WAPI_FAILED;
        }
    }

    return ret;
}

/**
 *  @brief software key set.
 */
int wai_cipher_setkey(td_u32 *keyslot_chn_num, td_u8 *key, td_u32 key_len, keyslot_engine keyslot_engine)
{
    int ret;
    td_handle klad_handle = 0;
    td_u32 salt[7] = {0x03020100, 0x07060504, 0x0b0a0908, 0x0f0e0d0c, 0x13121110, 0x14131514, 0x1b1a1918};
    uapi_drv_keyslot_type_t keyslot_type = UAPI_DRV_KEYSLOT_TYPE_MCIPHER;
    uapi_drv_klad_dest_t klad_type = UAPI_DRV_KLAD_DEST_MCIPHER;
    uapi_drv_klad_attr_t klad_attr = {
        .key_cfg = {
            .engine = UAPI_DRV_KLAD_ENGINE_AES,
            .decrypt_support = TRUE,
            .encrypt_support = TRUE
        },
        .key_sec_cfg = {
            .key_sec = FALSE,
            .master_only_enable = FALSE,
            .dest_buf_sec_support = FALSE,
            .dest_buf_non_sec_support = TRUE,
            .src_buf_sec_support = FALSE,
            .src_buf_non_sec_support = TRUE,
        }
    };
    uapi_drv_klad_clear_key_t clear_key = {
        .key = key,
        .key_length = key_len
    };
    uapi_drv_klad_effective_key_t content_key = {
        .salt = (td_u8 *)salt,
        .salt_length = sizeof(salt),
        .oneway = FALSE,
        .key_size = UAPI_DRV_KLAD_KEY_SIZE_128BIT,
        .key_parity = FALSE
    };

    klad_attr.key_cfg.engine = UAPI_DRV_KLAD_ENGINE_SM4;

    ret = uapi_drv_km_init();
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_km_init failed, ret is 0x%x\n", ret);
    }

    ret = uapi_drv_keyslot_create(keyslot_chn_num, keyslot_type);
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_keyslot_create failed, ret is 0x%x\n", ret);
    }
    osal_printk("create keyslot_num 0x%x\n", *keyslot_chn_num);

    ret = uapi_drv_klad_create(&klad_handle);
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_klad_create failed, ret is 0x%x\n", ret);
    }

    ret = uapi_drv_klad_attach(klad_handle, klad_type, *keyslot_chn_num);
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_klad_attach failed, ret is 0x%x\n", ret);
    }

    ret = uapi_drv_klad_set_attr(klad_handle, &klad_attr);
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_klad_set_attr failed, ret is 0x%x\n", ret);
    }

    ret = uapi_drv_klad_set_clear_key(klad_handle, &clear_key);
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_klad_set_clear_key failed, ret is 0x%x\n", ret);
    }

    ret = uapi_drv_klad_destroy(klad_handle);
    if (ret != ERRCODE_SUCC) {
        osal_printk("uapi_drv_klad_destroy failed, ret is 0x%x\n", ret);
    }
    return 0;
}

/* �ԳƼӽ��ܣ������ݽ��м��ܻ���� */
/**
 *  @brief Symmetric encryption and decryption
 */
static int wai_cipher_crypt(wai_operation_t operation, const unsigned char *iv, unsigned int ivlen,
    const unsigned char *key, unsigned int keylen, const unsigned char *input, unsigned int inlen,
    unsigned char *output, unsigned int *outlen)
{
    uint32_t symc_handle;
    uint32_t keyslot_handle;
    uapi_drv_cipher_symc_attr_t symc_attr = {0};
    uapi_drv_cipher_symc_ctrl_t symc_ctrl = {0};
    uapi_drv_cipher_buf_attr_t src_buf = {0};
    uapi_drv_cipher_buf_attr_t dst_buf = {0};
    uapi_drv_keyslot_type_t keyslot_type = {0};
    keyslot_engine engine_alg = KEYSLOT_ENGINE_AES;
    int ret = WAPI_FAILED;
    *outlen = WAI_NMK_SIZE;

    if (ivlen != UAPI_DRV_CIPHER_IV_LEN_IN_BYTES) {
        osal_printk(MSG_ERROR, "ivlen failed!");
        return WAPI_FAILED;
    }

    // ����������֮ǰ����������Կ
    engine_alg = KEYSLOT_ENGINE_SM4;
    ret = wai_cipher_setkey(&keyslot_handle, (unsigned char *)key, keylen, engine_alg);
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "cipher set key failed!");
        return ret;
    }

    // 1. ����uapi_drv_cipher_symc_init��ʼ��SYMCģ��
    ret = uapi_drv_cipher_symc_init();
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "symc init failed!");
        return WAPI_FAILED;
    }

    // 2. ��һ·symc������ȡsymc���������uapi_drv_cipher_symc_create�ӿ���ɡ�
    symc_attr.symc_alg = UAPI_DRV_CIPHER_SYMC_ALG_SM4;
    symc_attr.work_mode = UAPI_DRV_CIPHER_SYMC_WORK_MODE_OFB;
    symc_attr.is_long_term = TD_TRUE;
    ret = uapi_drv_cipher_symc_create(&symc_handle, &symc_attr);
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "create symc failed!");
        uapi_drv_cipher_symc_deinit();
        return WAPI_FAILED;
    }

    // 3. ��keyslotͨ����symc���������uapi_drv_cipher_symc_attach�ӿ����
    ret = uapi_drv_cipher_symc_attach(symc_handle, keyslot_handle);
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "Attach keyslot handle to encryption handle failed!");
        return WAPI_FAILED;
    }

    // 4. ����symc�㷨������Ϣ����������Կ�������ԣ���ʼ�������ӽ���λ����Կ���·�ʽ��GCM/CCMģʽ
    symc_ctrl.symc_alg = UAPI_DRV_CIPHER_SYMC_ALG_SM4;
    symc_ctrl.work_mode = UAPI_DRV_CIPHER_SYMC_WORK_MODE_OFB;
    symc_ctrl.symc_key_length = UAPI_DRV_CIPHER_SYMC_KEY_128BIT;
    symc_ctrl.symc_bit_width = UAPI_DRV_CIPHER_SYMC_BIT_WIDTH_128BIT;
    symc_ctrl.iv_length = ivlen;
    if (memcpy_s(symc_ctrl.iv, sizeof(symc_ctrl.iv), iv, ivlen) != EOK) {
        osal_printk(MSG_ERROR, "memcpy iv failed!");
        return WAPI_FAILED;
    }
    ret = uapi_drv_cipher_symc_set_config(symc_handle, &symc_ctrl);
    if (ret != ERRCODE_SUCC) {
        osal_printk(MSG_ERROR, "Set algorithm parameters failed!");
        uapi_drv_cipher_symc_destroy(symc_handle);
        return WAPI_FAILED;
    }

    // 5. �����ݽ��мӽ���
    src_buf.phys_addr = (uintptr_t)input;
    dst_buf.phys_addr = (uintptr_t)output;
    if (operation == WAI_ENCRYPT) {
        // ����
        ret = uapi_drv_cipher_symc_encrypt(symc_handle, &src_buf, &dst_buf, inlen);
        if (ret != WAPI_SUCCESS) {
            osal_printk(MSG_ERROR, "Symc encrypt failed!");
            uapi_drv_cipher_symc_destroy(symc_handle);
            return WAPI_FAILED;
        }
    } else if (operation == WAI_DECRYPT) {
        // ����
        ret = uapi_drv_cipher_symc_decrypt(symc_handle, &src_buf, &dst_buf, *outlen);
        if (ret != WAPI_SUCCESS) {
            osal_printk(MSG_ERROR, "Symc decrypt failed!");
            uapi_drv_cipher_symc_destroy(symc_handle);
            return WAPI_FAILED;
        }
    } else {
        osal_printk(MSG_ERROR, "operation type is error!");
        return WAPI_FAILED;
    }

    // 6. ���keyslotͨ����cipher���������uapi_drv_cipher_symc_detach�ӿ����
    ret = uapi_drv_cipher_symc_detach(symc_handle, keyslot_handle);
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "Symc detach failed!");
        return WAPI_FAILED;
    }

    // 7. ����symc���������uapi_drv_cipher_symc_destroy�ӿ����
    ret = uapi_drv_cipher_symc_destroy(symc_handle);
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "Symc destroy failed!");
        return WAPI_FAILED;
    }

    // 8. ����uapi_drv_cipher_symc_deinitȥ��ʼ��SYMCģ��
    ret = uapi_drv_cipher_symc_deinit();
    if (ret != WAPI_SUCCESS) {
        osal_printk(MSG_ERROR, "Symc deinit failed!");
        return WAPI_FAILED;
    }
    return WAPI_SUCCESS;
}

/**
 *  @brief SM4 encryption and decryption
 */
int wai_sm4_ofb_decrypt(const unsigned char *iv, unsigned int ivlen,
    const unsigned char *key, unsigned int keylen, const unsigned char *input, unsigned int inlen,
    unsigned char *output, unsigned int *outlen)
{
    return wai_cipher_crypt(WAI_DECRYPT, iv, ivlen,
        key, keylen, input, inlen, output, outlen);
}
